#!/usr/bin/env python
'''
This function extract the junctions information from sam file.  
It has two ways to use:  
(1) Given the .SAM file with junctions
supporting alignments, output the corresponding .bed file
@param args[0]: filename of the .SAM file
@param args[1]: output file name
@param args[2]: the target contig_name in the .SAM file
 
(2) Given the .SAM file with all junction-support mappings
of a contig, get all the mappings in the given range.
@param args[0]: filename of the .SAM file
@param args[1]: output_file name
@param args[2]: lower bound of the extracted alignments' left end.
@param args[3]: upper bound of the extracted alignments' left end.
@param args[4]: optional arg for shift of the index, default 0.
'''
__author__ = "Yangho Chen"
__email__ = "yanghoch@usc.edu"
__modification__ = "Created on Aug 16, 2009"

import sys
import os
import getopt
from search_a_large_sorted_file import SortedFileSearcher
from TophatOutFormat import JunctionLineInfo
from TophatOutFormat import SamLineInfo
from common import output_all_lines

def get_juncs_info_per_contig_from_bed( bed_filename, ofile_N, junc_no_threshold ):
    """ 
    Currently not used because tophat bed output has shift error
    (1) Given the .bed file (tophat's output), 
    extract junctions line of each contigs
    (genes) to a separate file.
    @param args[0]: filename of the .bed file
    """
    MIN_JUNCTIONS_NUM4OUTPUT = 2
    contigName = ""
    ifile = open( bed_filename, mode = "r" )
    line = ifile.readline() # ignore the first line header  
    # read until the first sequence header
    juncs = set() # set to collect junction info per contig 
    while 1:
        line = ifile.readline()
        if ( not line ) or line == '\n':
            #output the last 
            output_all_lines( ofile_N + contigName + ".bed", juncs )
            break;

        juncL = JunctionLineInfo( line )
        if ( juncL.gene != contigName ) :
            if( len( juncs ) >= MIN_JUNCTIONS_NUM4OUTPUT ):
                output_all_lines( contigName + ".bed", juncs )
            juncs.clear()
            contigName = juncL.gene
        juncs.add( line )
    ifile.close()

def juncsSupportAInSam2JuncsInBed( sam_file_name, ofileN, contig_name ):
    ''' 
    (2) Given the .SAM file which contains junctions supporting alignments.
    convert the junctions supporting alignments to junctions in BED format
    @param args[0]: filename of the .SAM file
    @param args[1]: contig_name of the .SAM file 
    '''
    def add_juncs( juncs_dict, juc_start, junc_end ):
        try:
            key = str( juc_start ) + '_' + str( junc_end )
            juncs_dict[key] = juncs_dict[key] + 1
        except KeyError:
            juncs_dict[key] = 1

    def output_juncs ( junc_dicts ):
        ofile = open( ofileN, mode = "w" )
        def output_juncs4strand( juncs_dict, strand ):
            for k in juncs_dict:
                js_e = k.split( '_' )
                juncStart = int( js_e[0] )
                juncEnd = int( js_e[1] )
                if juncStart >= 0 and juncEnd >= 0:
                    v = juncs_dict[k]
                    line = "%s\t%s\t%s\tJunc\t%s\t%c\n" % \
                        ( contig_name, js_e[0], js_e[1], v, strand )
                    ofile.write( line )

        output_juncs4strand( junc_dicts[0], '+' )
        output_juncs4strand( junc_dicts[1], '-' )
        ofile.close()

    ifile = open( sam_file_name, mode = "r" )
    juncs_dicts = [dict(), dict()]  # Forward and backward junctions containers
    while 1:
        line = ifile.readline()
        if ( not line ) or line == '\n':
            # output the junction file            
            output_juncs( juncs_dicts )
            break;
        sam_line = SamLineInfo( line )
        if( contig_name in sam_line.RNAME ):
            if( sam_line.get_junction_from_CIGAR() ):
                # add junction start, end and support reads count
                start = sam_line.junction_start
                end = sam_line.junction_end
                if sam_line.strand == '+':
                    add_juncs( juncs_dicts[0], start, end );
                else:
                    add_juncs( juncs_dicts[1], start, end );
    ifile.close()

def get_juncs_in_a_range_from_sam( sam_file_name, ofile_N, range, shift ):
    '''Given the .SAM file with all junction-support mapping
    of a contig, get all the mappings in the given range.
    @param args[0]: filename of the .SAM file
    @param args[1]: lower bound of the extracted alignments' left end.
    @param args[2]: upper bound of the extracted alignments' left end.'''
    posColId = 3;
    delimiter = '\t'
    s = SortedFileSearcher( sam_file_name );
    startPos = min( range )
    endPos = max( range )
    s.locatedFp( startPos, posColId, delimiter )
    records = s.get_lines_w_smaller_key( endPos, posColId , delimiter )
    print sam_file_name + ':' + str( len( records ) ) + " junction-support reads"
    ofile = open( ofile_N, "w" );
    for line in records:
        sam_line = SamLineInfo( line )
        ofile.write( sam_line.out_line( shift ) )
#        if( sam_line.get_junction_from_CIGAR() ):
#            if ( sam_line.junction_start >= range[0] ) and\
#               ( sam_line.junction_end <= range[1] ):
#                    ofile.write( sam_line.out_line( shift ) )
    ofile.close()

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error, msg:
        print msg
        sys.exit( 2 );
    #process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            print __doc__
            sys.exit();

    if( len( args ) >= 2 ):
        ifileN = args[0]
        ofileN = args[1]
        if( not os.path.exists( ifileN ) ):
            print ifileN, " doesn't exist."
            return False;

    if ( len( args ) == 3 ):
        targetContigN = args[2]
        juncsSupportAInSam2JuncsInBed( ifileN, ofileN, targetContigN )
    elif ( len( args ) == 4 ):
        chrRange = [int( args[2] ), int( args[3] )]
        get_juncs_in_a_range_from_sam( ifileN, ofileN, chrRange, 0 )
    elif ( len( args ) == 5 ):
        chrRange = [int( args[2] ), int( args[3] )]
        shift = int( args[4] )
        get_juncs_in_a_range_from_sam( ifileN, ofileN, chrRange, shift )
    else:
        usage()

if __name__ == '__main__':
    main()
